import json
import os

def labelme2yolo(json_dir, labels_dir, class_names):
    """
    将LabelMe标注的JSON文件转换为YOLO格式的TXT文件。
    :param json_dir: LabelMe标注的JSON文件所在目录。
    :param labels_dir: 转换后的TXT文件保存目录。
    :param class_names: 类别名称列表，用于将类别名称映射到类别索引。
    """
    if not os.path.exists(labels_dir):
        os.makedirs(labels_dir)
    
    for filename in os.listdir(json_dir):
        if filename.endswith('.json'):
            json_path = os.path.join(json_dir, filename)
            with open(json_path, 'r') as f:
                data = json.load(f)
            
            txt_filename = os.path.splitext(filename)[0] + '.txt'
            txt_path = os.path.join(labels_dir, txt_filename)
            
            with open(txt_path, 'w') as f:
                for shape in data['shapes']:
                    label = shape['label']
                    class_id = class_names.index(label)  # 将类别名称映射到索引
                    points = shape['points']
                    x = (points[0][0] + points[1][0]) / 2 / data['imageWidth']
                    y = (points[0][1] + points[1][1]) / 2 / data['imageHeight']
                    w = abs(points[1][0] - points[0][0]) / data['imageWidth']
                    h = abs(points[1][1] - points[0][1]) / data['imageHeight']
                    
                    f.write(f"{class_id} {x} {y} {w} {h}\n")

# 示例类别名称列表
class_names = ["lucy", "kina", "jmd", "ruihao", "haodi", "yry"]  # 根据实际情况替换类别名称

# 转换目录
json_dir = '/Users/v/Documents/proj/testdata/ai/release_data/img_mark'
labels_dir = '/Users/v/Documents/proj/testdata/ai/release_data/img_yolo_txt'

# 执行转换
labelme2yolo(json_dir, labels_dir, class_names)